package com.mjk.common.core.jdbc.ibatis;


import com.mjk.common.core.jdbc.datasource.DynamicDataSource;
import com.mjk.common.tools.logger.BizLogger;
import com.mjk.common.util.ReflectionUtils;
import com.mjk.common.util.StrUtils;
import org.apache.ibatis.transaction.Transaction;
import org.springframework.jdbc.CannotGetJdbcConnectionException;
import org.springframework.jdbc.datasource.DataSourceUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;


/**
 * 多数据源切换，支持事务
 * @author mjk
 *
 */
public class MultiDataSourceTransaction implements Transaction{
	
    private final DataSource dataSource;
 
    private Connection mainConnection;
 
    private String mainDatabaseIdentification;
 
    private final ConcurrentMap<String, Connection> otherConnectionMap;
 
    private boolean isConnectionTransactional;
 
    private boolean autoCommit;
    
    private final Map<Object, DataSource> resolved;
    
    @SuppressWarnings("unchecked")
	public MultiDataSourceTransaction(DataSource dataSource) {
        StrUtils.notNull(dataSource, "No DataSource specified");
        this.dataSource = dataSource;
        resolved = (Map<Object, DataSource>) ReflectionUtils.getFieldValue(dataSource, "resolvedDataSources");
        otherConnectionMap = new ConcurrentHashMap<>();
        mainDatabaseIdentification= DynamicDataSource.getDatabaseType();
    }
 
 
    /**
     * {@inheritDoc}
     */
    @Override
    public Connection getConnection() throws SQLException {
        String databaseIdentification = DynamicDataSource.getDatabaseType();
//        BizLogger.info("datasourcename :"+databaseIdentification);
        DataSource ds = resolved.get(databaseIdentification);
        if(null == ds) {
            ds = this.dataSource;
        }
        if (databaseIdentification.equals(mainDatabaseIdentification)) {
            if (mainConnection != null) {
                //do notiong
            }
            else {
                openMainConnection(dataSource);
                mainDatabaseIdentification =databaseIdentification;
            }
            return mainConnection;
        } else {
            if (!otherConnectionMap.containsKey(databaseIdentification)) {
                try {
                    Connection conn = ds.getConnection();
                    otherConnectionMap.put(databaseIdentification, conn);
                } catch (SQLException ex) {
                    throw new CannotGetJdbcConnectionException("Could not get JDBC Connection", ex);
                }
            }
            return otherConnectionMap.get(databaseIdentification);
        }
    }
 
 
    private void openMainConnection(DataSource dataSource) throws SQLException {
        this.mainConnection = DataSourceUtils.getConnection(dataSource);
        this.autoCommit = this.mainConnection.getAutoCommit();
        this.isConnectionTransactional = DataSourceUtils.isConnectionTransactional(this.mainConnection, dataSource);
         
        BizLogger.debug(
                    "JDBC Connection ["
                            + this.mainConnection
                            + "] will"
                            + (this.isConnectionTransactional ? " " : " not ")
                            + "be managed by Spring");        
    }
 
    /**
     * {@inheritDoc}
     */
    @Override
    public void commit() throws SQLException {
        if (this.mainConnection != null && !this.isConnectionTransactional && !this.autoCommit) {            
            BizLogger.debug("Committing JDBC Connection [" + this.mainConnection + "]");            
            this.mainConnection.commit();
            for (Connection connection : otherConnectionMap.values()) {
                connection.commit();
            }
        }
    }
 
    /**
     * {@inheritDoc}
     */
    @Override
    public void rollback() throws SQLException {
        if (this.mainConnection != null && !this.isConnectionTransactional && !this.autoCommit) {            
            BizLogger.debug("Rolling back JDBC Connection [" + this.mainConnection + "]");            
            this.mainConnection.rollback();
            for (Connection connection : otherConnectionMap.values()) {
                connection.rollback();
            }
        }
    }
 
    /**
     * {@inheritDoc}
     */
    @Override
    public void close() {
        DataSourceUtils.releaseConnection(this.mainConnection, this.dataSource);
        for (Connection connection : otherConnectionMap.values()) {
            DataSourceUtils.releaseConnection(connection, this.dataSource);
        }
    }
 
    @Override
    public Integer getTimeout() {
        return null;
    }

}
